package top.rainbowecho.ai.domain.model.robot.strategy.impl;

import org.springframework.stereotype.Component;
import top.rainbowecho.ai.domain.model.robot.AiRobot;
import top.rainbowecho.ai.domain.model.robot.strategy.SensitiveImageRecognizeStrategy;
import top.rainbowecho.ai.domain.model.robot.value.image.SensitiveImageListItem;
import top.rainbowecho.common.domain.shared.vo.SensitiveRecognizeResult;
import top.rainbowecho.common.util.CommonKey;
import top.rainbowecho.common.util.ExceptionMessageContent;

import java.util.*;

/**
 * @author rainbow
 * @since 2019/11/29 15:22
 */
@Component(AiRobot.TERRORISM_STRATEGY)
public class TerrorismImageRecognizeStrategy implements SensitiveImageRecognizeStrategy {
    /**
     * 暴恐标签集合
     */
    private Set<String> terrorismSet = new HashSet<>(Arrays.asList(CommonKey.ARMOREDCAR,
            CommonKey.TERRORISTS,
            CommonKey.OTHER_WEAPON,
            CommonKey.NORMALARMY,
            CommonKey.KNIFE,
            CommonKey.GUNS,
            CommonKey.CANNON,
            CommonKey.BLOOD));

    /**
     * 根据腾讯AI接口的图片鉴别接口响应结果进行识别
     *
     * @param reply 接口响应结果
     * @return 识别结果
     */
    @Override
    public SensitiveRecognizeResult recognize(List<SensitiveImageListItem> reply) {
        // 如果审核结果中最大置信度的标签在暴恐标签集合中，且置信度大于80则认为是暴恐图片

        Optional<SensitiveImageListItem> maxItem = reply.stream().max(Comparator.comparingInt(SensitiveImageListItem::getTag_confidence));
        SensitiveImageListItem maxConfidenceItem = maxItem.orElse(null);
        boolean contain = terrorismSet.contains(Objects.requireNonNull(maxConfidenceItem).getTag_name());
        boolean over = maxConfidenceItem.getTag_confidence() > 80;

        boolean terrorism = contain && over;
        String label = terrorism ? maxConfidenceItem.getTag_name() : ExceptionMessageContent.NORMAL;

        return new SensitiveRecognizeResult(terrorism, label);
    }
}
